{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import re\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import os\n",
    "import sys\n",
    "from tf2gpt.model import GPT\n",
    "from utils.story_util import Story,Stories\n",
    "from utils.progress_bar import ProgressBar\n",
    "from tensorboardX import SummaryWriter\n",
    "from tensorflow.keras.utils import multi_gpu_model\n",
    "import random\n",
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU'),\n",
       " PhysicalDevice(name='/physical_device:XLA_CPU:0', device_type='XLA_CPU'),\n",
       " PhysicalDevice(name='/physical_device:XLA_GPU:0', device_type='XLA_GPU'),\n",
       " PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.config.list_physical_devices()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Num GPUs Available:  1\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "print(\"Num GPUs Available: \", len(tf.config.experimental.list_physical_devices('GPU')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.keras.backend.set_floatx('float16')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#mirrored_strategy = tf.distribute.MirroredStrategy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "float16 <dtype: 'float16'> True\n"
     ]
    }
   ],
   "source": [
    "#with mirrored_strategy.scope():\n",
    "gpt = GPT(\n",
    "    vocab_size=30_000,\n",
    "    layer_size=32,\n",
    "    block_size=1024,\n",
    "    embedding_dropout=0.0,\n",
    "    embedding_size=2560,\n",
    "    num_attention_heads=32,\n",
    "    attention_dropout=0.0,\n",
    "    residual_dropout=0.0,\n",
    "    train_size=499\n",
    ")\n",
    "gpt.load_weights('./gpt_weight_pretrain/weight_fp16')\n",
    "\n",
    "#input_x = tf.keras.layers.Input((499,), dtype=tf.int32)\n",
    "#outputs = gpt_origin(input_x)\n",
    "\n",
    "#gpt = tf.keras.Model(inputs=input_x, outputs=outputs)\n",
    "#gpt = multi_gpu_model(gpt, gpus=8)\n",
    "\n",
    "print(tf.keras.backend.floatx(), tf.float16, tf.keras.backend.floatx() == tf.float16)\n",
    "if tf.keras.backend.floatx() == tf.float16:\n",
    "    for x in gpt.weights:\n",
    "        assert x.dtype == tf.float16\n",
    "\n",
    "\n",
    "gpt.compile(\n",
    "    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-5),  # Optimizer\n",
    "    # Loss function to minimize\n",
    "    loss=tf.keras.losses.SparseCategoricalCrossentropy(),\n",
    "    # List of metrics to monitor\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.gpt2_tokenizer import GPT2Tokenizer\n",
    "cbpe = GPT2Tokenizer(\n",
    "    'CPM-Generate/bpe_3w_new/vocab.json',\n",
    "    'CPM-Generate/bpe_3w_new/merges.txt',\n",
    "    model_file='CPM-Generate/bpe_3w_new/chinese_vocab.model')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Building prefix dict from the default dictionary ...\n",
      "Loading model from cache /tmp/jieba.cache\n",
      "Loading model cost 0.632 seconds.\n",
      "Prefix dict has been built successfully.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[837, 259, 497, 57, 8, 237]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ids = cbpe.encode('今天天气还行')\n",
    "ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "++++++++++++++++++++\n",
      "气 不错 不错行 \n",
      "------------------------------\n",
      "气 不错 不错行 ,\n",
      "------------------------------\n",
      "气 不错 不错行 , 我\n",
      "------------------------------\n",
      "气 不错 不错行 , 我 就\n",
      "------------------------------\n",
      "气 不错 不错行 , 我 就 想\n",
      "------------------------------\n",
      "气 不错 不错行 , 我 就 想着\n",
      "------------------------------\n",
      "气 不错 不错行 , 我 就 想着 去\n",
      "------------------------------\n",
      "气 不错 不错行 , 我 就 想着 去 看看\n",
      "------------------------------\n",
      "气 不错 不错行 , 我 就 想着 去 看看 \n",
      "------------------------------\n",
      "气 不错 不错行 , 我 就 想着 去 看看 ,\n",
      "------------------------------\n"
     ]
    }
   ],
   "source": [
    "def test_basic_logic():\n",
    "    ids = cbpe.encode('今天天气还行')\n",
    "    #print(ids)\n",
    "    print(\"+\" * 20)\n",
    "    for i in range(10):\n",
    "        output = gpt(tf.constant([ids]))\n",
    "        #print(output[0].shape)\n",
    "        nid = np.argmax(output[0, -1])\n",
    "        ids += [nid]\n",
    "        #print(i, cbpe.decode(ids))\n",
    "        #print(np.argmax(output[0],axis=-1))\n",
    "        print(cbpe.decode(np.argmax(output[0],axis=-1)))\n",
    "        print('-' * 30)\n",
    "test_basic_logic()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_learning_rate(learning_rate=6e-4,\n",
    "                      warmup_steps=20_0000,\n",
    "                      decay_steps=200_0000,\n",
    "                      alpha=0.0):\n",
    "    def decayed_learning_rate(step=1):\n",
    "        if step <= warmup_steps:\n",
    "            mult = step / float(warmup_steps)\n",
    "        else:\n",
    "            progress = (step - warmup_steps) / (decay_steps - warmup_steps)\n",
    "            mult = 0.5 * (1 + math.cos(math.pi * progress))\n",
    "            mult = max(0.1, mult)\n",
    "        return learning_rate * mult\n",
    "    return decayed_learning_rate"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "#stories = Stories(\"./labeled_data/advanture_translated/processed_translated_story.txt\").stories\n",
    "stories = Stories(\"./labeled_data/advanture_translated/processed_translated_story.txt\").stories\n",
    "#stories = stories[:50]\n",
    "data_folder = \"./labeled_data/\"\n",
    "txt_files = [(data_folder + i) for i in os.listdir(data_folder) if \"txt\" in i]\n",
    "#stories = stories[:10]\n",
    "stories = [Story(\"\",\"\").from_file(i) for i in txt_files]"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "valid_stories = Stories(\"./labeled_data/advanture_translated/processed_translated_story_valid.txt\").stories"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "len(stories),len(valid_stories)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_recoard = json.load(open('./labeled_data/advanture_translated/truncated_advanture_train.json'))\n",
    "valid_recoard = json.load(open('./labeled_data/advanture_translated/truncated_advanture_valid.json'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4500, 266)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(train_recoard),len(valid_recoard)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "def data_generator(stories, batch_size=4,sample_len=200,inf=False):\n",
    "    while True:\n",
    "        batch_data = []\n",
    "        tmp_stories = copy.copy(stories)\n",
    "        random.shuffle(tmp_stories)\n",
    "        for i,one_story in enumerate(tmp_stories):\n",
    "            story_content = one_story.to_dungeon_format()\n",
    "            story_content = story_content.replace(\"<start>\\n\",\"\")\n",
    "            story_content = story_content.replace(\"\\n<end>\",\"\")\n",
    "            story_content = story_content.replace(\"\\n<end>\",\"\")\n",
    "            story_content = story_content.replace(\" \",\"\")\n",
    "            ids = cbpe.encode(story_content)\n",
    "            while ids:\n",
    "                sample = ids[:sample_len]\n",
    "                ids = ids[sample_len:]\n",
    "                if len(sample) < sample_len:\n",
    "                    sample += [0 for i in range((sample_len - len(sample)))]\n",
    "                batch_data.append(sample)\n",
    "                if len(batch_data) >= batch_size:\n",
    "                    yield np.asarray(batch_data)\n",
    "                    batch_data = []\n",
    "        if not inf:\n",
    "            break\n",
    "            \n",
    "import copy\n",
    "def data_generator_content(texts, batch_size=4,sample_len=200,inf=False):\n",
    "    tmp_text = copy.copy(texts)\n",
    "    segments = []\n",
    "    while tmp_text:\n",
    "        story_content = tmp_text[:10000]\n",
    "        segments.append(story_content)\n",
    "        tmp_text = tmp_text[10000:]\n",
    "    while True:\n",
    "        \n",
    "        batch_data = []\n",
    "        random.shuffle(segments)\n",
    "        for one_story_context in segments:\n",
    "            ids = cbpe.encode(one_story_context)\n",
    "            while ids:\n",
    "                sample = ids[:sample_len]\n",
    "                ids = ids[sample_len:]\n",
    "                if len(sample) < sample_len:\n",
    "                    sample += [0 for i in range((sample_len - len(sample)))]\n",
    "                batch_data.append(sample)\n",
    "                if len(batch_data) >= batch_size:\n",
    "                    yield np.asarray(batch_data)\n",
    "                    batch_data = []\n",
    "        if not inf:\n",
    "            break\n",
    "            \n",
    "def turb_sentence(inputs):\n",
    "    outputs = []\n",
    "    for i in inputs:\n",
    "        if random.random() > 0.1:\n",
    "            outputs.append(i)\n",
    "    return outputs\n",
    "            \n",
    "def data_generator_json(stories, batch_size=1,inf=False):\n",
    "    while True:\n",
    "        tmp_stories = copy.copy(stories)\n",
    "        random.shuffle(tmp_stories)\n",
    "        for i,one_story in enumerate(tmp_stories):\n",
    "            previous_contents = one_story[\"previous\"]\n",
    "            previous_contents = turb_sentence(previous_contents)\n",
    "            action = one_story[\"action\"]\n",
    "            content = one_story[\"result\"]\n",
    "            \n",
    "            previous_contents = ''.join(previous_contents)\n",
    "            content = ''.join(content)\n",
    "            \n",
    "            ids_prev = cbpe.encode(previous_contents + \"\\n\")\n",
    "            ids_action = cbpe.encode(\">\" + action + \"\\n\")\n",
    "            ids_content = cbpe.encode(content)[:300]\n",
    "            \n",
    "            ids_prev = ids_prev + ids_action\n",
    "            \n",
    "            prev_context_len = int(random.uniform(100,500))\n",
    "            ids_prev = ids_prev[-prev_context_len:]\n",
    "            \n",
    "            overall_ids = ids_prev + ids_content\n",
    "            mask = [0 for i in ids_prev] + [1 for i in ids_content]\n",
    "            yield np.asarray([overall_ids]),np.asarray([mask])\n",
    "            \n",
    "        if not inf:\n",
    "            break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "valid_gen = data_generator_json(valid_recoard,inf=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "x,m = valid_gen.__next__()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((1, 311), (1, 311))"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.shape,m.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'起 。 你 觉得 有点 头晕 , 于是 你 把 自己 带到 黑暗 的 角落 休息 一下 。 您 访问 该 区域 角落 有点 暗 , 因为 自然光 几乎 无法 到达 那里 。 地板 上 似乎 也 有 一些 垃圾 。 你 进一步 移动 到 商店 的 后面 担心 周围 有 更 多 的 敌人 , 你 把 自己 拖 到 商店 的 后角 。 你 快 休息 你 无精打采 地 趴在 地上 , 试着 平稳 地 呼吸 。 要 放大 的 痛苦 众生 。 你 想 知道 这个 世界 发生 了 什么 , 是否 还有 更 多 的 敌对者 , 比如 攻击 你 的 那个 人 。 你 检查 后备箱 你 下车 , 走向 后备箱 。 你 走 回 商店 , 找到 几件 你 撕开 并 裹 在 胸前 的 衬衫 。 它 可以 帮助 您 感觉 稍微 好 一点 。 但 你 仍然 很 痛苦 。 你 弹出 后备箱 您 按 下 按钮 解锁 后备箱 。 它 缓缓 抬起 打开 。 你 发动 汽车 您 尝试 多次 启动 汽车 , 但 没有 任何 反应 。 你 想想 它 可能 需要 汽油 , 也许 他们 的 后备箱 里 有 备用 汽油 。 你 检查 地板 \\n > 你 忍住 了 咳嗽 \\n 当 你 看到 一个 女人 小心翼翼 地 走进 来时 , 你 把手 放在 嘴 上 。 她 看起来 并 没有 什么 害处 , 但 谁 知道 她 心里 在 想 什么 。'"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cbpe.decode(x[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "        1, 1, 1]])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1\n",
      "tl: 2.439453125 vl: 2.533203125 69.02 % [==================================>----------------] 3106/4500 \t used:3417s eta:1533 sss"
     ]
    }
   ],
   "source": [
    "writer = SummaryWriter('log/finetune')\n",
    "n_iter = 0\n",
    "train_loss = 100\n",
    "val_loss = 100\n",
    "for epoch in range(3):\n",
    "    print(f\"Epoch {epoch + 1}\")\n",
    "    pb = ProgressBar(len(train_recoard))\n",
    "    pb.startjob()\n",
    "    for x,m in data_generator_json(train_recoard):\n",
    "    #for x in data_generator_content(texts[528:],sample_len=500,batch_size=2):\n",
    "        n_iter += 1\n",
    "        ret = gpt.train_step(\n",
    "            (\n",
    "                tf.constant(x[:,:-1]),\n",
    "                tf.constant(x[:,1:])\n",
    "            ),\n",
    "            mask=m[:,1:],\n",
    "        )\n",
    "        writer.add_scalar(\"train/loss\", ret[\"loss\"].numpy(), n_iter)\n",
    "        train_loss = ret[\"loss\"].numpy()\n",
    "        if n_iter % 10 == 0:\n",
    "            valid_x,valid_m = valid_gen.__next__()\n",
    "            ret = gpt.eval_step(\n",
    "            (\n",
    "                tf.constant(valid_x[:,:-1]),\n",
    "                tf.constant(valid_x[:,1:])),\n",
    "                mask=valid_m[:,1:],\n",
    "            )\n",
    "            writer.add_scalar(\"test/loss\", ret[\"loss\"].numpy(), n_iter)\n",
    "            val_loss = ret[\"loss\"].numpy()\n",
    "        pb.info = f\"tl: {train_loss} vl: {val_loss}\"\n",
    "        pb.complete(1)\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gpt.save_weights('./gpt_weight_pretrain/tmp_weight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -l ./gpt_weight_pretrain/tmp_weight*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# rewrite fp16 model (weight only)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import re\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import os\n",
    "import sys\n",
    "from tf2gpt.model import GPT\n",
    "from utils.story_util import Story,Stories\n",
    "from utils.progress_bar import ProgressBar\n",
    "from tensorboardX import SummaryWriter\n",
    "from tensorflow.keras.utils import multi_gpu_model\n",
    "import random\n",
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Building prefix dict from the default dictionary ...\n",
      "Loading model from cache /tmp/jieba.cache\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "float16 <dtype: 'float16'> True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading model cost 0.559 seconds.\n",
      "Prefix dict has been built successfully.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "++++++++++++++++++++\n",
      "好 不错行 \n",
      "------------------------------\n",
      "好 不错行 ,\n",
      "------------------------------\n",
      "好 不错行 , \n",
      "------------------------------\n",
      "好 不错行 ,  \n",
      "------------------------------\n",
      "好 不错行 ,   \n",
      "------------------------------\n",
      "好 不错行 ,    \n",
      "------------------------------\n",
      "好 不错行 ,     \n",
      "------------------------------\n",
      "好 不错行 ,      \n",
      "------------------------------\n",
      "好 不错行 ,       \n",
      "------------------------------\n",
      "好 不错行 ,        \n",
      "------------------------------\n"
     ]
    }
   ],
   "source": [
    "tf.keras.backend.set_floatx('float16')\n",
    "#with mirrored_strategy.scope():\n",
    "gpt = GPT(\n",
    "    vocab_size=30_000,\n",
    "    layer_size=32,\n",
    "    block_size=1024,\n",
    "    embedding_dropout=0.0,\n",
    "    embedding_size=2560,\n",
    "    num_attention_heads=32,\n",
    "    attention_dropout=0.0,\n",
    "    residual_dropout=0.0,\n",
    "    train_size=499\n",
    ")\n",
    "gpt.load_weights('./gpt_weight_pretrain/tmp_weight')\n",
    "\n",
    "#input_x = tf.keras.layers.Input((499,), dtype=tf.int32)\n",
    "#outputs = gpt_origin(input_x)\n",
    "\n",
    "#gpt = tf.keras.Model(inputs=input_x, outputs=outputs)\n",
    "#gpt = multi_gpu_model(gpt, gpus=8)\n",
    "\n",
    "print(tf.keras.backend.floatx(), tf.float16, tf.keras.backend.floatx() == tf.float16)\n",
    "if tf.keras.backend.floatx() == tf.float16:\n",
    "    for x in gpt.weights:\n",
    "        assert x.dtype == tf.float16\n",
    "        \n",
    "from utils.gpt2_tokenizer import GPT2Tokenizer\n",
    "cbpe = GPT2Tokenizer(\n",
    "    'CPM-Generate/bpe_3w_new/vocab.json',\n",
    "    'CPM-Generate/bpe_3w_new/merges.txt',\n",
    "    model_file='CPM-Generate/bpe_3w_new/chinese_vocab.model')\n",
    "\n",
    "def test_basic_logic():\n",
    "    ids = cbpe.encode('今天天气还行')\n",
    "    #print(ids)\n",
    "    print(\"+\" * 20)\n",
    "    for i in range(10):\n",
    "        output = gpt(tf.constant([ids]))\n",
    "        #print(output[0].shape)\n",
    "        nid = np.argmax(output[0, -1])\n",
    "        ids += [nid]\n",
    "        #print(i, cbpe.decode(ids))\n",
    "        #print(np.argmax(output[0],axis=-1))\n",
    "        print(cbpe.decode(np.argmax(output[0],axis=-1)))\n",
    "        print('-' * 30)\n",
    "test_basic_logic()\n",
    "        \n",
    "gpt.save_weights('./gpt_weight_pretrain/tmp_weight_test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-rw-r--r-- 1 tiger tiger 5194222858 Jul 14 10:40 ./gpt_weight_pretrain/tmp_weight_test.data-00000-of-00001\n",
      "-rw-r--r-- 1 tiger tiger      34188 Jul 14 10:40 ./gpt_weight_pretrain/tmp_weight_test.index\n"
     ]
    }
   ],
   "source": [
    "!ls -l ./gpt_weight_pretrain/tmp_weight_test*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# save model (restart kernel here)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import re\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import os\n",
    "import sys\n",
    "from tf2gpt.model import GPT\n",
    "from utils.story_util import Story,Stories\n",
    "from utils.progress_bar import ProgressBar\n",
    "from tensorboardX import SummaryWriter\n",
    "from tensorflow.keras.utils import multi_gpu_model\n",
    "import random\n",
    "tf.keras.backend.set_floatx('float16')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.gpt2_tokenizer import GPT2Tokenizer\n",
    "cbpe = GPT2Tokenizer(\n",
    "    'CPM-Generate/bpe_3w_new/vocab.json',\n",
    "    'CPM-Generate/bpe_3w_new/merges.txt',\n",
    "    model_file='CPM-Generate/bpe_3w_new/chinese_vocab.model')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "float16 <dtype: 'float16'> True\n"
     ]
    }
   ],
   "source": [
    "#with mirrored_strategy.scope():\n",
    "gpt = GPT(\n",
    "    vocab_size=30_000,\n",
    "    layer_size=32,\n",
    "    block_size=1024,\n",
    "    embedding_dropout=0.0,\n",
    "    embedding_size=2560,\n",
    "    num_attention_heads=32,\n",
    "    attention_dropout=0.0,\n",
    "    residual_dropout=0.0\n",
    ")\n",
    "gpt.load_weights('./gpt_weight_pretrain/tmp_weight_test')\n",
    "\n",
    "#input_x = tf.keras.layers.Input((499,), dtype=tf.int32)\n",
    "#outputs = gpt_origin(input_x)\n",
    "\n",
    "#gpt = tf.keras.Model(inputs=input_x, outputs=outputs)\n",
    "#gpt = multi_gpu_model(gpt, gpus=8)\n",
    "\n",
    "print(tf.keras.backend.floatx(), tf.float16, tf.keras.backend.floatx() == tf.float16)\n",
    "if tf.keras.backend.floatx() == tf.float16:\n",
    "    for x in gpt.weights:\n",
    "        assert x.dtype == tf.float16\n",
    "\n",
    "\n",
    "#gpt.compile(\n",
    "#    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-5),  # Optimizer\n",
    "#    # Loss function to minimize\n",
    "#    loss=tf.keras.losses.SparseCategoricalCrossentropy(),\n",
    "#    # List of metrics to monitor\n",
    "#)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Building prefix dict from the default dictionary ...\n",
      "Loading model from cache /tmp/jieba.cache\n",
      "Loading model cost 0.581 seconds.\n",
      "Prefix dict has been built successfully.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 今天天气 不错 \n",
      "1 今天天气 不错 ,\n",
      "2 今天天气 不错 , 我\n",
      "3 今天天气 不错 , 我 想\n",
      "4 今天天气 不错 , 我 想 去\n",
      "5 今天天气 不错 , 我 想 去 看看\n",
      "6 今天天气 不错 , 我 想 去 看看 \n",
      "7 今天天气 不错 , 我 想 去 看看 。\n",
      "8 今天天气 不错 , 我 想 去 看看 。 \n",
      "9 今天天气 不错 , 我 想 去 看看 。 ”\n"
     ]
    }
   ],
   "source": [
    "ids = cbpe.encode('今天天气不错')\n",
    "\n",
    "for i in range(10):\n",
    "    output = gpt(tf.constant([ids]))\n",
    "    nid = np.argmax(output[0, -1])\n",
    "    ids += [nid]\n",
    "    print(i, cbpe.decode(ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def batch_gather(a, b):\n",
    "    return tf.gather(a, b, batch_dims=1)\n",
    "\n",
    "def penalize_used(logits, output,penalize):\n",
    "    # I want to change the indices of logits wherever the index is found in output\n",
    "    change_tensor = tf.zeros_like(logits, dtype=logits.dtype)\n",
    "    unique = tf.unique(output[0])[0]\n",
    "    ones = tf.ones_like(unique, dtype=unique.dtype)\n",
    "    indices = tf.expand_dims(unique, 1)\n",
    "    updates = tf.scatter_nd(indices, ones, [logits.shape[1]])\n",
    "    bool_tensor = tf.expand_dims(tf.cast(updates, tf.bool), 0)\n",
    "    return tf.compat.v1.where(bool_tensor, logits * penalize, logits)\n",
    "\n",
    "def top_p_sample(logits, inputs, num_samples=1, p=0.9,k=40,temperature=0.4,penalize=0.4):\n",
    "    logits = penalize_used(logits,inputs,penalize)\n",
    "    batch_size, vocab_size = logits.shape\n",
    "    probs = tf.nn.softmax(logits / temperature, axis=-1)\n",
    "    # [batch_size, vocab_perm]\n",
    "    indices = tf.argsort(probs, direction='DESCENDING')\n",
    "    cumulative_probabilities = tf.math.cumsum(batch_gather(probs, indices), axis=-1, exclusive=False)\n",
    "    top_k_prob = cumulative_probabilities[0][k]\n",
    "\n",
    "    # find the top pth index to cut off. careful we don't want to cutoff everything!\n",
    "    # result will be [batch_size, vocab_perm]\n",
    "    p_expanded = p if isinstance(p, float) else p#[:, None]\n",
    "    exclude_mask = tf.logical_not(\n",
    "        tf.logical_or(\n",
    "            tf.logical_and(\n",
    "                cumulative_probabilities < p_expanded, \n",
    "                cumulative_probabilities < top_k_prob\n",
    "            ),\n",
    "            tf.range(vocab_size)[None] < 1\n",
    "        )\n",
    "    )\n",
    "\n",
    "    # OPTION A - sample in the sorted space, then unsort.\n",
    "    logits_to_use = batch_gather(logits, indices) - tf.cast(exclude_mask, tf.float16) * 1e4\n",
    "    sample_perm = tf.random.categorical(logits=logits_to_use, num_samples=num_samples)\n",
    "    sample = batch_gather(indices, sample_perm)\n",
    "\n",
    "    return tf.cast(sample, tf.int64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def serve(inputs):\n",
    "    return gpt(inputs, kv_cache=None, use_cache=True)\n",
    "\n",
    "\n",
    "@tf.function\n",
    "def serve_cache(inputs, kv_cache):\n",
    "    return gpt(inputs, kv_cache=kv_cache, use_cache=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "serve_concrete = serve.get_concrete_function(\n",
    "    tf.TensorSpec(shape=[None, None], dtype=tf.int64, name=\"inp\")\n",
    ")\n",
    "\n",
    "layer_size = 32\n",
    "attention_head = 32\n",
    "embedding_size = 2560\n",
    "\n",
    "serve_cache_concrete = serve_cache.get_concrete_function(\n",
    "    tf.TensorSpec(shape=[None, None], dtype=tf.int64, name=\"inp\"),\n",
    "    tf.TensorSpec(shape=[\n",
    "        layer_size, None, 2, attention_head,\n",
    "        None, embedding_size // attention_head\n",
    "    ], dtype=tf.float16, name=\"kv_cache\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "r = serve_concrete(\n",
    "    tf.constant([[1]], tf.int64)\n",
    ")\n",
    "r2 = serve_cache_concrete(\n",
    "    tf.constant([[1]], tf.int64),\n",
    "    r[1]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def sample(initial_inputs, length, p,k,temperature,penalize, num_samples=1):\n",
    "    layer_size = 32\n",
    "    embedding_size = 2560\n",
    "    attention_head = 32\n",
    "\n",
    "    i = tf.constant(0, dtype=tf.int64)\n",
    "    initial_logits, kv_cache = serve(initial_inputs)\n",
    "    inputs = top_p_sample(initial_logits[:, -1, :],initial_inputs, p=p,k=k,temperature=temperature,penalize=penalize)\n",
    "    stores = tf.concat([initial_inputs, inputs], axis=1)\n",
    "\n",
    "    def _cond(i, inputs, kv_cache, stores):\n",
    "        return i < length\n",
    "\n",
    "    def _body(i, inputs, kv_cache, stores):\n",
    "        new_logits, new_kv_cache = serve_cache(inputs, kv_cache)\n",
    "        \n",
    "        new_inputs = top_p_sample(new_logits[:, -1, :],inputs, p=p,k=k,temperature=temperature,penalize=penalize)\n",
    "        new_stores = tf.concat([stores, new_inputs], axis=-1)\n",
    "        new_kv_cache = tf.concat([\n",
    "            kv_cache,\n",
    "            new_kv_cache\n",
    "        ], axis=-2)\n",
    "        new_i = i + 1\n",
    "        return [new_i, new_inputs, new_kv_cache, new_stores]\n",
    "\n",
    "    result = tf.while_loop(\n",
    "        _cond, _body,\n",
    "        loop_vars=[i, inputs, kv_cache, stores],\n",
    "        shape_invariants=[\n",
    "            tf.TensorShape(None),\n",
    "            tf.TensorShape([None, None]),\n",
    "            tf.TensorShape([\n",
    "                layer_size, None, 2,\n",
    "                attention_head, None,\n",
    "                embedding_size // attention_head\n",
    "            ]),\n",
    "            tf.TensorShape([\n",
    "                None, None\n",
    "            ])\n",
    "        ]\n",
    "    )\n",
    "    return result[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /usr/local/lib/python3.7/dist-packages/tensorflow/python/training/tracking/tracking.py:111: Model.state_updates (from tensorflow.python.keras.engine.training) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "This property should not be used in TensorFlow 2.0, as updates are applied automatically.\n",
      "WARNING:tensorflow:From /usr/local/lib/python3.7/dist-packages/tensorflow/python/training/tracking/tracking.py:111: Layer.updates (from tensorflow.python.keras.engine.base_layer) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "This property should not be used in TensorFlow 2.0, as updates are applied automatically.\n",
      "INFO:tensorflow:Assets written to: ./result_weights/tmp_weight/assets\n"
     ]
    }
   ],
   "source": [
    "gpt.save('./result_weights/tmp_weight', include_optimizer=False, signatures={\n",
    "    'serving_default': sample.get_concrete_function(\n",
    "        tf.TensorSpec(shape=[None, None], dtype=tf.int64, name=\"inp\"),\n",
    "        tf.TensorSpec(shape=[None,], dtype=tf.int64, name=\"length\"),\n",
    "        tf.TensorSpec(shape=[], dtype=tf.float16, name=\"p\"),\n",
    "        tf.TensorSpec(shape=[], dtype=tf.int64, name=\"k\"),\n",
    "        tf.TensorSpec(shape=[], dtype=tf.float16, name=\"temperature\"),\n",
    "        tf.TensorSpec(shape=[], dtype=tf.float16, name=\"penalize\")\n",
    "    )\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[ 837  259  497  788    8    9  148   16  755   91 1126  375    8   12\n",
      "     8]], shape=(1, 15), dtype=int64)\n",
      "今天天气 不错 , 所以 我 决定 去 海边 。 \n"
     ]
    }
   ],
   "source": [
    "ids = cbpe.encode('今天天气不错')\n",
    "\n",
    "ret = sample(\n",
    "    tf.constant([ids], dtype=tf.int64), # inp\n",
    "    tf.constant(10, dtype=tf.int64), #length\n",
    "    tf.constant(0.9, dtype=tf.float16), #p\n",
    "    tf.constant(40, dtype=tf.int64), #k\n",
    "    tf.constant(0.4, dtype=tf.float16), # temperature\n",
    "    tf.constant(0.4, dtype=tf.float16) # penalize\n",
    ")\n",
    "print(ret)\n",
    "print(cbpe.decode(ret.numpy().tolist()[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 21660\n",
      "drwxr-xr-x 2 tiger tiger     4096 Jul  5 15:21 assets\n",
      "-rw-r--r-- 1 tiger tiger 22168232 Jul 14 10:50 saved_model.pb\n",
      "drwxr-xr-x 3 tiger tiger     4096 Jul 14 10:50 variables\n"
     ]
    }
   ],
   "source": [
    "!ls -l ./result_weights/tmp_weight"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# read deploy model (you should restart kernel and go from here)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import re\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import os\n",
    "import sys\n",
    "from tf2gpt.model import GPT\n",
    "from utils.story_util import Story,Stories\n",
    "from utils.progress_bar import ProgressBar\n",
    "from tensorflow.keras.utils import multi_gpu_model\n",
    "import random\n",
    "from utils.gpt2_tokenizer import GPT2Tokenizer\n",
    "import tensorflow_hub as hub\n",
    "cbpe = GPT2Tokenizer(\n",
    "    'CPM-Generate/bpe_3w_new/vocab.json',\n",
    "    'CPM-Generate/bpe_3w_new/merges.txt',\n",
    "    model_file='CPM-Generate/bpe_3w_new/chinese_vocab.model')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "gpt = hub.load('./result_weights/tmp_weight/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_gpt(tokenizer, gpt, sentence, number=1, length=20, p=0.9,k=40,temperature=0.4,penalize=0.85):\n",
    "    inputs = tf.constant([tokenizer.encode(sentence)] * number, dtype=tf.int64)\n",
    "    length = tf.constant(length, dtype=tf.int64)\n",
    "    p = tf.constant(p, dtype=tf.float16)\n",
    "    k = tf.constant(k, dtype=tf.int64)\n",
    "    temperature = tf.constant(temperature, dtype=tf.float16)\n",
    "    penalize = tf.constant(penalize, dtype=tf.float16)\n",
    "    ret = gpt.signatures['serving_default'](\n",
    "        inp=inputs, \n",
    "        length=length, \n",
    "        p=p,\n",
    "        k=k,\n",
    "        temperature=temperature,\n",
    "        penalize=penalize)['output_0']\n",
    "    return [\n",
    "        tokenizer.decode(s[len(inputs[0]):]).replace(' ', '')\n",
    "        for s in ret.numpy()\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Building prefix dict from the default dictionary ...\n",
      "Loading model from cache /tmp/jieba.cache\n",
      "Loading model cost 0.558 seconds.\n",
      "Prefix dict has been built successfully.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "了吧。你可以看看他的回答。<eod>有人说FastestisthestupidthingFastestisthestupidthing有人说FastestisthestupidthingFastestisthestupidthingisFastestisthestupidthingFastestisthestupidthingFastestisthestupidthingFastestisthestupidthing.........FastestisthestupidthingFastestisthestupidthingFastestisthestupidthing.........FastestisthestupidthingFastestisthestupidthingFastestisthestupidthing.........Fastestisthestupidthing.........FastestisthestupidthingFastestisthestupidthingFastestisthestupidthing.........Fastestisthestupidthing...\n",
      "--------------------\n"
     ]
    }
   ],
   "source": [
    "q = f'''双十一'''\n",
    "ret = sample_gpt(cbpe, gpt, q, 1, 500, p=0.9,k=16,temperature=1,penalize=0.75)\n",
    "for x in ret:\n",
    "    print(x)\n",
    "    print('-' * 20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Story():\n",
    "    def __init__(self,beginning,story_max_len=200,context_len=12):\n",
    "        self.story = [beginning]\n",
    "        self.story_max_len = story_max_len\n",
    "        self.context_len = context_len\n",
    "        \n",
    "    def response_quality_ok(self,response):\n",
    "        if len(response) < 20:\n",
    "            return False\n",
    "        return True\n",
    "    \n",
    "    def action(self,action):\n",
    "        action_str = \"> 你\" + action\n",
    "        q = \"\".join(self.story[-self.context_len:]).replace(\"\\n\",'') + \"\\n\" + action_str + \"\\n\"\n",
    "        q = q[-600:]#.replace(\"\\n\",\"\")\n",
    "        self.story.append(\"你\" + action)\n",
    "        response = \"\"\n",
    "        for i in range(3):\n",
    "            response = sample_gpt(cbpe, gpt, q, 1, 150,p=0.9,k=16,temperature=1,penalize=0.75)[0]\n",
    "            response = '。'.join(response.split(\"。\")[:-1]) + \"。\"\n",
    "            if(self.response_quality_ok(response) == True):\n",
    "                break\n",
    "            else:\n",
    "                print(\"response quality check failed,generating another one\")\n",
    "        \n",
    "        #responsesp = response.split(\"。\")\n",
    "        #if(len(responsesp) > 2):\n",
    "        #    response = ''.join(responsesp[:-1])\n",
    "        #else:\n",
    "        #    response = response.split(\">\")[0]\n",
    "        #response = \"。\".join([:-1])\n",
    "        #response = response.split(\">\")[0].strip()\n",
    "        #response = \"\\n\".join(response.split(\"\\n\")[:-1])\n",
    "        self.story.append(response)\n",
    "        \n",
    "    def interactive(self):\n",
    "        print(\"\\n\".join(self.story))\n",
    "        while True:\n",
    "            action = input(\"> 你\")\n",
    "            self.action(action)\n",
    "            print(self.story[-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "begin = f'''你在树林里冒险，指不定会从哪里蹦出来一些奇怪的东西，你握紧手上的手枪，希望这次冒险能够找到一些值钱的东西'''\n",
    "story = Story(begin)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "你在树林里冒险，指不定会从哪里蹦出来一些奇怪的东西，你握紧手上的手枪，希望这次冒险能够找到一些值钱的东西\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "> 你 把手枪上膛\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "枪身立刻就被拉下来,你把弹匣从枪口推出来,把枪扔到地上。“现在,我们走吧。”你说。“好的,先生。”老人说。“那我们出发吧。”你说。“好的,先生。”老人说。你把手枪放进自己的口袋里,然后跟着老人继续旅行。“我们去哪里?”老人问。“去Geralt村。\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "> 你 走去村子里\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "很快,你就会到达Geralt村。村子里有个院子,院子里有个漂亮的石头房子。你走进房子,发现门廊上坐着一个很漂亮的女人,她的脸颊红润,看起来很漂亮。“嘿,丽贝卡,”你喊道,“我是来帮你的。”“谢谢,亲爱的。”“什么?”“我叫你丽贝卡,我是你的妻子。\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "> 你 挽着她散步\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "当你和丽贝卡交谈时,你开始觉得很舒服,你想着你在这里找到了很多东西,你终于开始享受这个地方了。你的心情越来越好,你想要的东西就好像越来越近了。“那么,亲爱的,我们要去哪里?”“很明显,我们该走了。我们很快就会到村子里的Geralt村。”你说。“好的,那么,我们到那里吧。\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "> 你 在村子里找个地方玩\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "你带着老人和丽贝卡到村子里最好的房子去,并开始找一个最好的。你很快就找到了。它就在村子里,你会找到它的。你带着老人和丽贝卡来到村子的边缘,然后丽贝卡就走在了你的前面。\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "> 你 走进房子\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "您带着老人和丽贝卡进入房子,然后在房子里等。你走到房子的门口,你把门厅里的椅子推到一边,让老人坐下。你走到房子的另一边,你把丽贝卡带到房子另一边,并把老人从椅子上抱下来。你把老人放下,丽贝卡就坐在你的膝盖上。“怎么了?我有什么要说的?”你问。“我们需要你和我们。我有我的妻子和你的妻子。\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "> 你 赞美妻子的美貌\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "“你真棒,亲爱的。”丽贝卡说。“哦,是的,但她太年轻了。”“我会照顾她的。您看到她的样子了吗?”“是的,她看起来很漂亮。”“你会看到它,你看。她的眼睛好看吗?”“是的,很亮。”“好的,”你说。“那么,我们走吧。你可以和我们一起去Geralt村。\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "> 你 走去村里的集市\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "现在,你去集市。这个集市在村子里,你把老人和丽贝卡留在你的领子上,然后你带着老人丽贝卡到外面的广场。你在那里看到了一个大的摊位,上面摆着面包和葡萄酒。你把老人和丽贝卡带到那里,你把老人从椅子上抱下来,然后你把她放下,开始和她交谈。\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "> 你 说“你觉得这个集市怎么样？”\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "“你认为怎么样?我想去Geralt村,我们可以去村里的大集市。你看到了一个大的布店,那里有很多布。我们可以在那里做生意。你看到了一个大的水果摊,你可以买你需要的任何水果。我们可以做到了。”“是的,”你说。“我认为我们会做到的。我们可以做到的。现在走吧,我们去Geralt村。我想看看我们会做到什么。我认为我们已经做到了。\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "Interrupted by user",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-7-c73048edb922>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mstory\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minteractive\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m<ipython-input-5-b8be70995cce>\u001b[0m in \u001b[0;36minteractive\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     37\u001b[0m         \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"\\n\"\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjoin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstory\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     38\u001b[0m         \u001b[0;32mwhile\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 39\u001b[0;31m             \u001b[0maction\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"> 你\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     40\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maction\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0maction\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     41\u001b[0m             \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstory\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/ipykernel/kernelbase.py\u001b[0m in \u001b[0;36mraw_input\u001b[0;34m(self, prompt)\u001b[0m\n\u001b[1;32m    849\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_parent_ident\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    850\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_parent_header\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 851\u001b[0;31m             \u001b[0mpassword\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    852\u001b[0m         )\n\u001b[1;32m    853\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/ipykernel/kernelbase.py\u001b[0m in \u001b[0;36m_input_request\u001b[0;34m(self, prompt, ident, parent, password)\u001b[0m\n\u001b[1;32m    890\u001b[0m             \u001b[0;32mexcept\u001b[0m \u001b[0mKeyboardInterrupt\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    891\u001b[0m                 \u001b[0;31m# re-raise KeyboardInterrupt, to truncate traceback\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 892\u001b[0;31m                 \u001b[0;32mraise\u001b[0m \u001b[0mKeyboardInterrupt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Interrupted by user\"\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    893\u001b[0m             \u001b[0;32mexcept\u001b[0m \u001b[0mException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    894\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlog\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwarning\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Invalid Message:\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mexc_info\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: Interrupted by user"
     ]
    }
   ],
   "source": [
    "story.interactive()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
